package cloud.seri.gateway.security.oauth2

import cloud.seri.gateway.web.rest.errors.InvalidPasswordException
import io.github.jhipster.security.PersistentTokenCache
import org.slf4j.LoggerFactory
import org.springframework.http.ResponseEntity
import org.springframework.security.oauth2.common.OAuth2AccessToken
import org.springframework.web.client.HttpClientErrorException

import javax.servlet.http.Cookie
import javax.servlet.http.HttpServletRequest
import javax.servlet.http.HttpServletResponse

/**
 * Manages authentication cases for OAuth2 updating the cookies holding access and refresh tokens accordingly.
 *
 * It can authenticate users, refresh the token cookies should they expire and log users out.
 */
class OAuth2AuthenticationService(
    /**
     * Used to contact the OAuth2 token endpoint.
     */
    private val authorizationClient: OAuth2TokenEndpointClient,
    /**
     * Helps us with cookie handling.
     */
    private val cookieHelper: OAuth2CookieHelper) {

    private val log = LoggerFactory.getLogger(OAuth2AuthenticationService::class.java)

    /**
     * Caches Refresh grant results for a refresh token value so we can reuse them.
     * This avoids hammering UAA in case of several multi-threaded requests arriving in parallel.
     */
    private val recentlyRefreshed: PersistentTokenCache<OAuth2Cookies>

    init {
        recentlyRefreshed = PersistentTokenCache(REFRESH_TOKEN_VALIDITY_MILLIS)
    }

    /**
     * Authenticate the user by username and password.
     *
     * @param request  the request coming from the client.
     * @param response the response going back to the server.
     * @param params   the params holding the username, password and rememberMe.
     * @return the [OAuth2AccessToken] as a [ResponseEntity]. Will return `OK (200)`, if successful.
     * If the UAA cannot authenticate the user, the status code returned by UAA will be returned.
     */
    fun authenticate(request: HttpServletRequest, response: HttpServletResponse,
                     params: Map<String, String>): ResponseEntity<OAuth2AccessToken> {
        try {
            val username = params["username"]
            val password = params["password"]
            val rememberMe = java.lang.Boolean.valueOf(params["rememberMe"])
            val accessToken = authorizationClient.sendPasswordGrant(username!!, password!!)
            val cookies = OAuth2Cookies()
            cookieHelper.createCookies(request, accessToken, rememberMe, cookies)
            cookies.addCookiesTo(response)
            if (log.isDebugEnabled) {
                log.debug("successfully authenticated user {}", params["username"])
            }
            return ResponseEntity.ok(accessToken)
        } catch (ex: HttpClientErrorException) {
            log.error("failed to get OAuth2 tokens from UAA", ex)
            throw InvalidPasswordException()
        }
    }

    /**
     * Try to refresh the access token using the refresh token provided as cookie.
     * Note that browsers typically send multiple requests in parallel which means the access token
     * will be expired on multiple threads. We don't want to send multiple requests to UAA though,
     * so we need to cache results for a certain duration and synchronize threads to avoid sending
     * multiple requests in parallel.
     *
     * @param request       the request potentially holding the refresh token.
     * @param response      the response setting the new cookies (if refresh was successful).
     * @param refreshCookie the refresh token cookie. Must not be null.
     * @return the new servlet request containing the updated cookies for relaying downstream.
     */
    fun refreshToken(request: HttpServletRequest, response: HttpServletResponse, refreshCookie: Cookie): HttpServletRequest {
        //check if non-remember-me session has expired
        if (cookieHelper.isSessionExpired(refreshCookie)) {
            log.info("session has expired due to inactivity")
            logout(request, response)       //logout to clear cookies in browser
            return stripTokens(request)            //don't include cookies downstream
        }
        val cookies = getCachedCookies(refreshCookie.value)
        synchronized(cookies) {
            //check if we have a result from another thread already
            if (cookies.accessTokenCookie == null) {            //no, we are first!
                //send a refresh_token grant to UAA, getting new tokens
                val refreshCookieValue = OAuth2CookieHelper.getRefreshTokenValue(refreshCookie)
                val accessToken = authorizationClient.sendRefreshGrant(refreshCookieValue)
                val rememberMe = OAuth2CookieHelper.isRememberMe(refreshCookie)
                cookieHelper.createCookies(request, accessToken, rememberMe, cookies)
                //add cookies to response to update browser
                cookies.addCookiesTo(response)
            } else {
                log.debug("reusing cached refresh_token grant")
            }
            //replace cookies in original request with new ones
            val requestCookies = CookieCollection(*request.cookies)
            cookies.accessTokenCookie?.apply { requestCookies.add(this) }
            cookies.refreshTokenCookie?.apply { requestCookies.add(this) }
            return CookiesHttpServletRequestWrapper(request, requestCookies.toArray())
        }
    }

    /**
     * Get the result from the cache in a thread-safe manner.
     *
     * @param refreshTokenValue the refresh token for which we want the results.
     * @return a RefreshGrantResult for that token. This will either be empty, if we are the first one to do the
     * request, or contain some results already, if another thread already handled the grant for us.
     */
    private fun getCachedCookies(refreshTokenValue: String): OAuth2Cookies {
        synchronized(recentlyRefreshed) {
            var ctx: OAuth2Cookies? = recentlyRefreshed.get(refreshTokenValue)
            if (ctx == null) {
                ctx = OAuth2Cookies()
                recentlyRefreshed.put(refreshTokenValue, ctx)
            }
            return ctx
        }
    }

    /**
     * Logs the user out by clearing all cookies.
     *
     * @param httpServletRequest  the request containing the Cookies.
     * @param httpServletResponse the response used to clear them.
     */
    fun logout(httpServletRequest: HttpServletRequest, httpServletResponse: HttpServletResponse) {
        cookieHelper.clearCookies(httpServletRequest, httpServletResponse)
    }

    /**
     * Strips token cookies preventing them from being used further down the chain.
     * For example, the OAuth2 client won't checked them and they won't be relayed to other services.
     *
     * @param httpServletRequest the incoming request.
     * @return the request to replace it with which has the tokens stripped.
     */
    fun stripTokens(httpServletRequest: HttpServletRequest): HttpServletRequest {
        val cookies = cookieHelper.stripCookies(httpServletRequest.cookies)
        return CookiesHttpServletRequestWrapper(httpServletRequest, cookies)
    }

    companion object {
        /**
         * Number of milliseconds to cache refresh token grants so we don't have to repeat them in case of parallel requests.
         */
        private const val REFRESH_TOKEN_VALIDITY_MILLIS = 10000L
    }
}
